"""
python p2_solution_simple.py
"""

import json
import os
from pulp import LpProblem, LpMinimize, LpVariable, lpSum, LpBinary, value

# 文件路径
warehouse_file = '../fujian/fujian3/origin_data/warehouse.json'
inventory_file = '../fujian/fujian3/data_from_p1/all_average_inventory.json'
sales_file     = '../fujian/fujian3/data_from_p1/all_average_sales.json'
output_file = '../fujian/fujian3/solution_simple/1c1p_solution.json'

# 数据读取
with open(warehouse_file, 'r', encoding='utf-8') as f:
    warehouses = json.load(f)
with open(inventory_file, 'r', encoding='utf-8') as f:
    inventories = json.load(f)
with open(sales_file, 'r', encoding='utf-8') as f:
    sales = json.load(f)

# 获取 M 和 N
M = len(inventories)
N = len(warehouses)
print(f"Total number of categories (M): {M}")
print(f"Total number of warehouses (N): {N}")

# 构建库存和销量的字典，方便后续查找
inventory_dict = {str(item['category_id']): item['average_inventory'] for item in inventories}
sales_dict = {str(item['category_id']): item['average_sales'] for item in sales}
warehouse_dict = {
    warehouse['warehouse_id']: {
        'max_inventory': warehouse['max_inventory'],
        'max_sales': warehouse['max_sales'],
        'daily_cost': warehouse['daily_cost']
    } for warehouse in warehouses
}

# 初始化问题
problem = LpProblem("Warehouse_Assignment_Problem", LpMinimize)

# 创建决策变量 x_{ij}
x = LpVariable.dicts("x", [(i, j) for i in inventory_dict.keys() for j in warehouse_dict.keys()], cat=LpBinary)

# 目标函数：最小化仓库总租用成本
# 如果仓库 j 被使用 (x_{ij} = 1 for some i)，则计算 daily_cost[j]
problem += lpSum([x[i, j] * warehouse_dict[j]['daily_cost'] for i in inventory_dict.keys() for j in warehouse_dict.keys()]), "Minimize_Cost"

# 约束条件 1：每个品类仅分配到一个仓库
for i in inventory_dict.keys():
    problem += lpSum([x[i, j] for j in warehouse_dict.keys()]) == 1, f"One_Warehouse_for_Category_{i}"

# 约束条件 2：每个仓库的库存量不超过最大库存限制
for j in warehouse_dict.keys():
    problem += lpSum([inventory_dict[i] * x[i, j] for i in inventory_dict.keys()]) <= warehouse_dict[j]['max_inventory'], f"Inventory_Limit_for_Warehouse_{j}"

# 约束条件 3：每个仓库的出货量不超过最大销量限制
for j in warehouse_dict.keys():
    problem += lpSum([sales_dict[i] * x[i, j] for i in sales_dict.keys()]) <= warehouse_dict[j]['max_sales'], f"Sales_Limit_for_Warehouse_{j}"

# 求解
problem.solve()

# 输出结果到 JSON
solution = []
for i in inventory_dict.keys():
    for j in warehouse_dict.keys():
        if value(x[i, j]) == 1:
            solution.append({"category_id": i, "warehouse_id": j})
            break  # 每个 category 只对应一个 warehouse

# 保存到 JSON 文件
with open(output_file, 'w', encoding='utf-8') as f:
    json.dump(solution, f, indent=4)

print(f"Solution saved to {output_file}")
